#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UCEIS集成模型 - 完全兼容现有inference_server.py
"""

import torch
import torch.nn as nn
import torchvision.models as models
from typing import Dict, Any
import numpy as np
from pathlib import Path


class UCEISScoreModel(nn.Module):
    """
    UCEIS评分模型 - 与现有推理服务器完全兼容
    输出格式与inference_server.py期望的一致
    """

    def __init__(self, config: Dict[str, Any]):
        super().__init__()

        self.config = config
        self.num_classes = config.get('num_classes', 9)  # 使用9分类匹配原始期望

        # 设置架构参数
        backbone_type = config.get('backbone_type', 'resnet50')
        pretrained = config.get('pretrained', True)
        feature_dims = config.get('feature_dims', [64, 128, 256, 512])

        # 创建backbone
        if backbone_type == 'resnet50':
            self.backbone = models.resnet50(pretrained=pretrained)
            num_ftrs = self.backbone.fc.in_features
            self.backbone.fc = nn.Identity()  # 移除原始分类器
        else:
            raise ValueError(f"Unsupported backbone: {backbone_type}")

        # UCEIS组件分类头
        self.vascular_classifier = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 3)  # 血管模式: 0-2分
        )

        self.bleeding_classifier = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 4)  # 出血: 0-3分
        )

        self.erosion_classifier = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 4)  # 糜烂/溃疡: 0-3分
        )

        # 总体分类器（作为后备）
        self.total_classifier = nn.Sequential(
            nn.Linear(num_ftrs, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(256, 8)  # 总体UCEIS: 0-7分
        )

        # 加载预训练权重（如果可用）
        self.model_path = config.get('model_path', None)
        if self.model_path:
            self.load_pretrained_weights()

    def load_pretrained_weights(self):
        """从预训练模型加载权重"""
        try:
            if self.model_path and Path(self.model_path).exists():
                print(f"从 {self.model_path} 加载预训练权重...")
                checkpoint = torch.load(self.model_path, map_location='cpu')

                if 'model_state_dict' in checkpoint:
                    state_dict = checkpoint['model_state_dict']
                else:
                    state_dict = checkpoint

                # 尝试加载兼容的权重
                self.load_partial_weights(state_dict)
                print("权重加载成功")
            else:
                print(f"模型文件不存在: {self.model_path}")
        except Exception as e:
            print(f"加载权重失败: {e}")

    def load_partial_weights(self, state_dict):
        """部分加载兼容的权重"""
        model_dict = self.state_dict()

        # 过滤并匹配权重
        compatible_dict = {}
        for k, v in state_dict.items():
            if 'backbone' in k:
                # 尝试匹配backbone权重
                new_key = k.replace('image_branch.backbone', 'backbone')
                if new_key in model_dict and v.shape == model_dict[new_key].shape:
                    compatible_dict[new_key] = v
            elif 'classifier' in k and 'total_classifier' in k:
                # 尝试匹配分类器权重
                new_key = k.replace('classifier', 'total_classifier')
                if new_key in model_dict and v.shape == model_dict[new_key].shape:
                    compatible_dict[new_key] = v

        # 更新模型权重
        model_dict.update(compatible_dict)
        self.load_state_dict(model_dict)

        print(f"加载了 {len(compatible_dict)} 个兼容的权重参数")

    def forward(self, x):
        """
        前向传播

        Args:
            x: 输入图像张量 (B, 3, H, W)

        Returns:
            dict: 包含各个组件的logits
        """
        # 提取特征
        features = self.backbone(x)  # (B, num_ftrs)

        # 各组件分类
        vascular_logits = self.vascular_classifier(features)  # (B, 3)
        bleeding_logits = self.bleeding_classifier(features)  # (B, 4)
        erosion_logits = self.erosion_classifier(features)  # (B, 4)
        total_logits = self.total_classifier(features)  # (B, 8)

        return {
            'vascular_logits': vascular_logits,
            'bleeding_logits': bleeding_logits,
            'erosion_logits': erosion_logits,
            'total_logits': total_logits  # 总体分类作为后备
        }

    def predict_comprehensive_score(self, x):
        """预测综合UCEIS评分"""
        with torch.no_grad():
            outputs = self.forward(x)

            # 获取各组件预测
            vascular_score = torch.argmax(outputs['vascular_logits'], dim=1)
            bleeding_score = torch.argmax(outputs['bleeding_logits'], dim=1)
            erosion_score = torch.argmax(outputs['erosion_logits'], dim=1)
            total_score = torch.argmax(outputs['total_logits'], dim=1)

            # 计算总分（主要方法）
            uceis_score = vascular_score + bleeding_score + erosion_score
            uceis_score = torch.clamp(uceis_score, 0, 8)  # 限制在0-8分

            # 后备：使用总体分类
            backup_score = total_score

            return {
                'uceis_score': uceis_score[0].item(),
                'backup_score': backup_score[0].item(),
                'vascular_score': vascular_score[0].item(),
                'bleeding_score': bleeding_score[0].item(),
                'erosion_score': erosion_score[0].item(),
                'outputs': outputs
            }

    def get_model_info(self):
        """获取模型信息"""
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)

        return {
            'model_type': 'UCEISScoreModel',
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'config': self.config,
            'components': {
                'vascular_pattern': {'classes': 3, 'description': '血管模式 (0-2分)'},
                'bleeding': {'classes': 4, 'description': '出血 (0-3分)'},
                'erosions_ulcers': {'classes': 4, 'description': '糜烂/溃疡 (0-3分)'},
                'total_classification': {'classes': 8, 'description': '总体UCEIS (0-7分)'}
            }
        }


def create_integrated_uceis_model(model_path: str = None) -> UCEISScoreModel:
    """
    创建集成的UCEIS模型

    Args:
        model_path: 预训练模型路径

    Returns:
        配置好的UCEISScoreModel实例
    """
    config = {
        'backbone_type': 'resnet50',
        'pretrained': True,
        'num_classes': 9,
        'feature_dims': [64, 128, 256, 512],
        'model_path': model_path
    }

    model = UCEISScoreModel(config)
    return model


def test_integrated_model():
    """测试集成模型"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 创建模型
    model_path = "outputs/simple_test/best_model.pth"
    model = create_integrated_uceis_model(model_path).to(device)

    # 获取模型信息
    info = model.get_model_info()
    print(f"模型信息:")
    for key, value in info.items():
        print(f"  {key}: {value}")

    # 测试前向传播
    batch_size = 2
    test_input = torch.randn(batch_size, 3, 224, 224).to(device)

    print(f"\n测试输入尺寸: {test_input.shape}")

    with torch.no_grad():
        # 基础前向传播
        outputs = model(test_input)
        print(f"输出格式: {list(outputs.keys())}")
        print(f"血管logits: {outputs['vascular_logits'].shape}")
        print(f"出血logits: {outputs['bleeding_logits'].shape}")
        print(f"糜烂logits: {outputs['erosion_logits'].shape}")
        print(f"总体logits: {outputs['total_logits'].shape}")

        # 综合预测
        prediction = model.predict_comprehensive_score(test_input)
        print(f"\n预测结果:")
        for key, value in prediction.items():
            if key != 'outputs':
                print(f"  {key}: {value}")

    return model


if __name__ == "__main__":
    test_integrated_model()